import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import xgboost as xgb
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import classification_report, accuracy_score
from sklearn.model_selection import train_test_split
import os
import glob
import chardet
import re
from datetime import datetime


# 1. 数据获取与预处理
def load_and_preprocess_data(file_path):
    try:
        # 检测文件编码
        with open(file_path, 'rb') as f:
            raw_data = f.read()
            result = chardet.detect(raw_data)
            encoding = result['encoding'] if result['encoding'] else 'utf-8'

        # 加载数据
        data = pd.read_csv(file_path, encoding=encoding)

        # 转换日期时间格式（处理可能的格式不一致）
        def fix_date(date_str):
            # 处理日期格式问题：2025-3-10 11:02 -> 2025-03-10 11:02
            date_str = re.sub(r'(\d{4})-(\d{1})-', r'\1-0\2-', date_str)
            date_str = re.sub(r' (\d{1}):', r' 0\1:', date_str)
            return date_str

        if 'trade_date' in data.columns:
            data['trade_date'] = data['trade_date'].apply(fix_date)
            data['trade_date'] = pd.to_datetime(data['trade_date'], errors='coerce')
            data = data.dropna(subset=['trade_date'])
            data = data.sort_values('trade_date').reset_index(drop=True)
        else:
            print(f"文件 {file_path} 缺少 'trade_date' 列")
            return None

        # 数据清洗 - 确保所有数值列都是数字类型
        numeric_cols = ['open', 'high', 'low', 'close', 'vol', 'amount', 'pos']
        for col in numeric_cols:
            if col in data.columns:
                data[col] = pd.to_numeric(data[col], errors='coerce')
            else:
                print(f"文件 {file_path} 缺少列: {col}")
                return None

        # 删除任何可能包含NaN值的行
        data = data.dropna()

        # 添加时间特征
        data['hour'] = data['trade_date'].dt.hour
        data['minute'] = data['trade_date'].dt.minute
        data['day_of_week'] = data['trade_date'].dt.dayofweek

        return data
    except Exception as e:
        print(f"处理文件 {file_path} 时出错: {str(e)}")
        return None


# 2. 特征工程
def feature_engineering(data, lookahead=5):
    # 计算价格变化率
    data['price_change'] = data['close'].pct_change()
    data['vol_change'] = data['vol'].pct_change()

    # 计算移动平均
    windows = [5, 10, 20]
    for w in windows:
        data[f'ma{w}'] = data['close'].rolling(window=w).mean()

    # 计算波动率
    data['volatility'] = data['close'].rolling(window=20).std()

    # 计算RSI
    delta = data['close'].diff()
    gain = (delta.where(delta > 0, 0)).fillna(0)
    loss = (-delta.where(delta < 0, 0)).fillna(0)

    avg_gain = gain.rolling(window=14).mean()
    avg_loss = loss.rolling(window=14).mean()
    rs = avg_gain / avg_loss
    data['RSI'] = 100 - (100 / (1 + rs))

    # 计算MACD
    exp12 = data['close'].ewm(span=12, adjust=False).mean()
    exp26 = data['close'].ewm(span=26, adjust=False).mean()
    data['MACD'] = exp12 - exp26
    data['Signal'] = data['MACD'].ewm(span=9, adjust=False).mean()

    # 创建标签 - 未来价格方向
    future_prices = data['close'].shift(-lookahead)
    price_change = (future_prices - data['close']) / data['close']

    # 设置阈值创建标签 - 修复标签格式问题
    data['label'] = 1  # 默认观望 -> 映射为1
    data.loc[price_change > 0.002, 'label'] = 2  # 做多 -> 映射为2
    data.loc[price_change < -0.002, 'label'] = 0  # 做空 -> 映射为0

    # 删除包含NaN的行
    data = data.dropna()

    return data


# 3. 模型构建与训练
def build_and_train_model(data):
    # 划分特征和标签
    features = data.drop(['trade_date', 'label'], axis=1)
    labels = data['label']

    # 检查标签类别数量
    unique_labels = labels.unique()
    if len(unique_labels) < 3:
        print(f"警告: 数据只包含 {len(unique_labels)} 个标签类别: {unique_labels}，无法训练三分类模型")
        return None, None

    # 划分训练集和测试集（70%-30%）
    train_size = int(len(data) * 0.7)
    X_train = features.iloc[:train_size]
    y_train = labels.iloc[:train_size]
    X_test = features.iloc[train_size:]
    y_test = labels.iloc[train_size:]

    # 标准化特征
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)

    # 创建并训练XGBoost模型
    model = xgb.XGBClassifier(
        objective='multi:softmax',
        num_class=3,
        n_estimators=300,
        max_depth=7,
        learning_rate=0.05,
        subsample=0.8,
        colsample_bytree=0.8,
        random_state=42,
        eval_metric='mlogloss'
    )

    model.fit(X_train, y_train)

    # 评估模型
    y_pred = model.predict(X_test)
    print("\n模型评估报告:")
    print(classification_report(y_test, y_pred))
    print(f"测试集准确率: {accuracy_score(y_test, y_pred):.4f}")

    return model, scaler


# 4. 交易策略与回测
class TradingSimulator:
    def __init__(self, initial_capital=100000):
        self.initial_capital = initial_capital
        self.capital = initial_capital
        self.position = 0  # 0:无持仓, 1:多头, -1:空头
        self.entry_price = 0
        self.trades = []
        self.equity = [initial_capital]
        self.current_equity = initial_capital
        self.trade_count = 0
        self.win_count = 0
        self.max_drawdown = 0
        self.max_equity = initial_capital

    def execute_trade(self, signal, current_price, current_time):
        # 交易成本（0.1%手续费）
        trade_cost = self.capital * 0.001

        # 平仓逻辑
        if self.position != 0:
            # 多头平仓
            if self.position == 1:
                pnl = (current_price - self.entry_price) * (self.capital * 0.1 / self.entry_price)
                self.capital = self.capital * 0.9 + (self.capital * 0.1 + pnl) - trade_cost
                trade_type = 'long_close'
                is_win = pnl > 0

            # 空头平仓
            elif self.position == -1:
                pnl = (self.entry_price - current_price) * (self.capital * 0.1 / self.entry_price)
                self.capital = self.capital * 0.9 + (self.capital * 0.1 + pnl) - trade_cost
                trade_type = 'short_close'
                is_win = pnl > 0

            # 记录交易
            self.trades.append({
                'type': trade_type,
                'price': current_price,
                'time': current_time,
                'pnl': pnl,
                'is_win': is_win,
                'equity': self.capital
            })

            # 更新胜率统计
            self.trade_count += 1
            if is_win:
                self.win_count += 1

            self.position = 0
            self.entry_price = 0

        # 开仓逻辑
        if signal != 1 and self.position == 0 and self.capital > 1000:  # 不为观望信号
            # 10%仓位
            trade_amount = self.capital * 0.1

            if signal == 2:  # 做多
                self.position = 1
                self.entry_price = current_price
                trade_type = 'long_open'

            elif signal == 0:  # 做空
                self.position = -1
                self.entry_price = current_price
                trade_type = 'short_open'

            # 记录交易
            self.trades.append({
                'type': trade_type,
                'price': current_price,
                'time': current_time,
                'pnl': 0,
                'is_win': None,
                'equity': self.capital
            })

            # 扣除交易成本
            self.capital -= trade_cost

    def update_equity(self, current_price):
        # 更新当前权益
        if self.position == 1:  # 多头持仓
            position_value = (self.current_equity * 0.1) * (current_price / self.entry_price)
            self.current_equity = self.current_equity * 0.9 + position_value

        elif self.position == -1:  # 空头持仓
            position_value = (self.current_equity * 0.1) * (self.entry_price / current_price)
            self.current_equity = self.current_equity * 0.9 + position_value

        else:  # 无持仓
            self.current_equity = self.capital

        # 更新最大回撤
        if self.current_equity > self.max_equity:
            self.max_equity = self.current_equity

        drawdown = (self.max_equity - self.current_equity) / self.max_equity
        if drawdown > self.max_drawdown:
            self.max_drawdown = drawdown

        self.equity.append(self.current_equity)

    def get_performance_metrics(self):
        win_rate = self.win_count / self.trade_count if self.trade_count > 0 else 0
        total_return = (self.current_equity - self.initial_capital) / self.initial_capital * 100  # 转换为百分比

        return {
            'win_rate': win_rate,
            'total_return': total_return,
            'max_drawdown': self.max_drawdown * 100,  # 转换为百分比
            'trade_count': self.trade_count
        }


# 5. 回测与可视化
def backtest_and_visualize(data, model, scaler, file_name):
    # 准备回测数据
    features = data.drop(['trade_date', 'label'], axis=1)
    scaled_features = scaler.transform(features)

    # 预测信号
    data['signal'] = model.predict(scaled_features)

    # 初始化交易模拟器
    simulator = TradingSimulator()

    # 运行回测
    for i, row in data.iterrows():
        current_price = row['close']
        current_time = row['trade_date']
        signal = row['signal']

        simulator.execute_trade(signal, current_price, current_time)
        simulator.update_equity(current_price)

    # 获取性能指标
    metrics = simulator.get_performance_metrics()
    trades_df = pd.DataFrame(simulator.trades)

    # 可视化结果
    plt.figure(figsize=(15, 12))

    # 价格和交易信号图
    plt.subplot(3, 1, 1)
    plt.plot(data['trade_date'], data['close'], label='Price', color='blue', alpha=0.7)

    # 标记交易信号
    if not trades_df.empty:
        long_opens = trades_df[trades_df['type'] == 'long_open']
        long_closes = trades_df[trades_df['type'] == 'long_close']
        short_opens = trades_df[trades_df['type'] == 'short_open']
        short_closes = trades_df[trades_df['type'] == 'short_close']

        plt.scatter(long_opens['time'], long_opens['price'],
                    color='green', marker='^', s=100, label='Long Open')
        plt.scatter(long_closes['time'], long_closes['price'],
                    color='blue', marker='v', s=100, label='Long Close')
        plt.scatter(short_opens['time'], short_opens['price'],
                    color='red', marker='v', s=100, label='Short Open')
        plt.scatter(short_closes['time'], short_closes['price'],
                    color='orange', marker='^', s=100, label='Short Close')

    plt.title(f'Price Chart with Trading Signals - {file_name}')
    plt.xlabel('Time')
    plt.ylabel('Price')
    plt.legend()
    plt.grid(True)

    # 资金曲线图
    plt.subplot(3, 1, 2)
    equity_curve = simulator.equity[:len(data)]
    plt.plot(data['trade_date'], equity_curve,
             label='Equity Curve', color='purple')
    plt.title('Equity Curve')
    plt.xlabel('Time')
    plt.ylabel('Equity')
    plt.grid(True)

    # 收益分布图
    plt.subplot(3, 1, 3)
    if not trades_df.empty and 'is_win' in trades_df.columns:
        trades_df = trades_df[trades_df['is_win'].notnull()]
        win_trades = trades_df[trades_df['is_win'] == True]
        loss_trades = trades_df[trades_df['is_win'] == False]

        if not win_trades.empty:
            plt.bar(win_trades['time'], win_trades['pnl'],
                    color='green', label='Profitable Trades')
        if not loss_trades.empty:
            plt.bar(loss_trades['time'], loss_trades['pnl'],
                    color='red', label='Loss Trades')

    plt.title('Trade Profit/Loss Distribution')
    plt.xlabel('Time')
    plt.ylabel('P/L')
    plt.legend()
    plt.grid(True)

    plt.tight_layout()

    # 为每个文件创建唯一的结果文件名
    base_name = os.path.splitext(file_name)[0]
    plt.savefig(f'results_{base_name}.png', dpi=300)

    return metrics, trades_df, plt


# 主函数
def main():
    # 获取当前文件所在目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    data_dir = os.path.join(current_dir, 'data')

    print("=" * 50)
    print("现货交易策略模型 - 多文件处理")
    print("=" * 50)

    # 检查data目录是否存在
    if not os.path.exists(data_dir):
        print(f"错误：目录不存在 - {data_dir}")
        print("请确认以下内容：")
        print("1. 在脚本同级创建了data目录")
        print("2. 数据文件已放入data目录")
        return

    # 获取所有CSV文件
    csv_files = glob.glob(os.path.join(data_dir, '*.csv'))
    if not csv_files:
        print(f"错误：{data_dir} 目录中没有CSV文件")
        return

    print(f"找到 {len(csv_files)} 个CSV文件:")
    for i, file_path in enumerate(csv_files):
        print(f"{i + 1}. {os.path.basename(file_path)}")

    # 创建结果目录
    results_dir = os.path.join(current_dir, 'results')
    if not os.path.exists(results_dir):
        os.makedirs(results_dir)

    # 汇总结果
    summary_results = []

    # 处理每个CSV文件
    for file_path in csv_files:
        file_name = os.path.basename(file_path)
        print(f"\n{'=' * 30} 开始处理文件: {file_name} {'=' * 30}")

        # 1. 加载并预处理数据
        print("\n[步骤1/5] 加载并预处理数据...")
        data = load_and_preprocess_data(file_path)
        if data is None or data.empty:
            print(f"⚠ 文件 {file_name} 加载失败或数据为空，跳过处理")
            continue

        print(f"数据加载完成，共 {len(data)} 行")

        # 2. 特征工程
        print("\n[步骤2/5] 进行特征工程...")
        try:
            data = feature_engineering(data)
            print(f"特征工程完成，剩余数据行数: {len(data)}")
        except Exception as e:
            print(f"特征工程出错: {str(e)}")
            continue

        # 3. 构建并训练模型
        print("\n[步骤3/5] 构建并训练模型...")
        try:
            model, scaler = build_and_train_model(data)
            if model is None or scaler is None:
                print(f"⚠ 模型训练失败，跳过文件 {file_name}")
                continue
        except Exception as e:
            print(f"模型训练出错: {str(e)}")
            continue

        # 4. 回测与可视化
        print("\n[步骤4/5] 进行回测与可视化...")
        try:
            metrics, trades_df, plot = backtest_and_visualize(data, model, scaler, file_name)

            # 保存图表
            base_name = os.path.splitext(file_name)[0]
            plot.savefig(os.path.join(results_dir, f'results_{base_name}.png'), dpi=300)
            plt.close(plot.gcf())  # 关闭图表释放内存
        except Exception as e:
            print(f"回测与可视化出错: {str(e)}")
            continue

        # 5. 保存结果
        print("\n[步骤5/5] 保存结果...")
        trades_df.to_csv(os.path.join(results_dir, f'trading_records_{base_name}.csv'), index=False)

        # 打印性能指标
        print("\n策略性能指标:")
        print(f"胜率: {metrics['win_rate']:.2%}")
        print(f"总收益率: {metrics['total_return']:.2f}%")
        print(f"最大回撤: {metrics['max_drawdown']:.2f}%")
        print(f"交易次数: {metrics['trade_count']}")

        # 添加到汇总
        summary_results.append({
            'file': file_name,
            'win_rate': metrics['win_rate'],
            'total_return': metrics['total_return'],
            'max_drawdown': metrics['max_drawdown'],
            'trade_count': metrics['trade_count']
        })

        print(f"\n✔ 文件 {file_name} 处理完成")
        print(f"✔ 交易记录已保存到 results/trading_records_{base_name}.csv")
        print(f"✔ 可视化结果已保存到 results/results_{base_name}.png")

    # 汇总所有文件结果
    if summary_results:
        print("\n\n" + "=" * 50)
        print("所有文件处理完成 - 汇总结果")
        print("=" * 50)

        # 创建汇总表格
        summary_df = pd.DataFrame(summary_results)

        # 计算平均值
        avg_win_rate = summary_df['win_rate'].mean()
        avg_return = summary_df['total_return'].mean()
        avg_drawdown = summary_df['max_drawdown'].mean()
        total_trades = summary_df['trade_count'].sum()

        summary_df.loc['平均值'] = [
            '平均值',
            avg_win_rate,
            avg_return,
            avg_drawdown,
            total_trades
        ]

        # 保存汇总结果
        summary_path = os.path.join(results_dir, 'summary_results.csv')
        summary_df.to_csv(summary_path, index=False)
        print(f"\n✔ 汇总结果已保存到 results/summary_results.csv")

        # 打印汇总
        print("\n汇总性能指标:")
        print(f"平均胜率: {avg_win_rate:.2%}")
        print(f"平均总收益率: {avg_return:.2f}%")
        print(f"平均最大回撤: {avg_drawdown:.2f}%")
        print(f"总交易次数: {total_trades}")


if __name__ == "__main__":
    main()